# 将鼾声数据集 分成可以训练的数据集
# 需要三种数据  （音频，气压） -> 音频标签

import os
import librosa
import pandas as pd
from pydub import AudioSegment


# 0 各个文件的路径
wav_label_path = "E:\\data\\鼾声多模态\\止鼾枕标签\\止鼾枕清洗后的数据_out"
wav_files_path = "E:\\data\\鼾声多模态\\止鼾枕清洗后的数据"
output_files_path = "E:\\data\\鼾声多模态\\止鼾枕清洗后的数据2"

# 这里是一些下面用到的方法


# 将 00：01：01 转换为 61秒
def time_to_seconds(time_str):
    """将时间字符串转换为秒数"""
    h, m, s = map(int, time_str.split(":"))
    return h * 3600 + m * 60 + s


# 将音频分段
def split_audio(input_file, output_dir, wav_label_list: list, press_data_list: list, group_name="segment", segment_length=5000):
    """
    将音频文件分割为指定长度的小片段。

    :param input_file: 输入音频文件路径。
    :param output_dir: 输出片段保存的目录。
    :param segment_length: 每个片段的长度，单位为毫秒。
    """
    # 确保输出目录存在
    os.makedirs(output_dir, exist_ok=True)
    
    # 加载音频文件
    audio = AudioSegment.from_file(input_file)
    
    # 获取音频时长（毫秒）
    total_length = len(audio)
    
    # 分割音频
    for i in range(0, total_length, segment_length):
        start_time = i
        end_time = min(i + segment_length, total_length)
        segment = audio[start_time:end_time]

        # 读取wav_label 文件判断是否发生打鼾
        if wav_label_list[i//segment_length + 1] == 0:
            segment_filename = os.path.join(output_dir, "0", f"{group_name}_{i // segment_length + 1}_0_.wav")
            press_filename = os.path.join(output_dir, "0", f"{group_name}_{i // segment_length + 1}_0_.txt")
            with open(press_filename, 'w') as out_file:
                try:
                    out_file.write("\n".join(map(str, press_data_list[i//segment_length])))
                    segment.export(segment_filename, format="wav")
                    print(f"已保存: {press_filename}")
                except:
                    # 保存报错的文件
                    with open("E:\\data\\鼾声多模态\\止鼾枕清洗后的数据2\\error.txt", "a") as error_file:
                        error_file.write(press_filename)
                        error_file.write("\n")
        else:
            print(i//segment_length + 1)
            segment_filename = os.path.join(output_dir, "1", f"{group_name}_{i // segment_length + 1}_1_.wav")
            press_filename = os.path.join(output_dir, "1", f"{group_name}_{i // segment_length + 1}_1_.txt")
            with open(press_filename, 'w') as out_file:
                try:
                    out_file.write("\n".join(map(str, press_data_list[i//segment_length])))
                    segment.export(segment_filename, format="wav")
                    print(f"已保存: {press_filename}")
                except:
                     # 保存报错的文件
                    with open("E:\\data\\鼾声多模态\\止鼾枕清洗后的数据2\\error.txt", "a") as error_file:
                        error_file.write(press_filename)
        
        # 保存片段
        # segment_filename = os.path.join(output_dir, f"segment_{i // segment_length + 1}.wav")
        # segment.export(segment_filename, format="wav")
        print(f"已保存片段: {segment_filename}")
        # break


def generate_labels(audio_duration, annotations, interval=5, threshold=0.5):
    """
    根据置信度标记音频的每个时间片段。
    
    :param audio_duration: 音频总时长（秒）。
    :param annotations: 标签数据，包含开始时间、结束时间和置信度。
    :param interval: 时间片段长度（秒）。
    :param threshold: 置信度阈值。
    :return: 每个时间片段的标签列表。
    """
    num_segments = audio_duration // interval
    labels = [0] * (num_segments + 1) 
    
    for _, row in annotations.iterrows():
        start = time_to_seconds(row["开始时间"])
        end = time_to_seconds(row["结束时间"])
        confidence = row["置信度"]
        
        # print(confidence)
        if confidence > threshold:
            start_segment = start // interval
            end_segment = end // interval
            for i in range(start_segment, end_segment + 1):
                # print(i)
                labels[i] = 1
        # break
                
    return labels


# 处理文本 气压值
def process_log(input_file, output_dir, interval=5):
    """
    从日志文件中提取每行最后 6 个数的最大值，并将每 interval 秒的数据保存到单独的文件中。
    
    :param input_file: 输入日志文件路径。
    :param output_dir: 输出目录。
    :param interval: 分组时间间隔，单位为秒。
    """
    # 确保输出目录存在
    os.makedirs(output_dir, exist_ok=True)
    
    # 存储提取的数据
    grouped_data = []
    current_group = []
    start_time = None
    
    # 读取日志文件
    with open(input_file, 'r') as file:
        for line in file:
            line = line.strip()
            if not line:
                continue  # 跳过空行
            
            # 提取时间戳和数据部分
            try:
                timestamp, data = line.split("]:")
                timestamp = timestamp.strip("[")  # 提取干净的时间戳
                data_numbers = list(map(int, data.split()[3:]))  # 提取数字部分，从第4列开始
            except ValueError:
                print(f"解析失败的行: {line}")
                continue
            
            # 获取最后 6 个数字中的最大值
            if len(data_numbers) >= 6:
                max_value = max(data_numbers[-6:])
            else:
                print(f"数据不足6列: {line}")
                continue
            
            # 计算当前时间点（以秒为单位）
            time_parts = timestamp.split(":")
            time_in_seconds = int(time_parts[0].split(" ")[1]) * 3600 + int(time_parts[1]) * 60 + int(time_parts[2])
            
            # 初始化 start_time
            if start_time is None:
                start_time = time_in_seconds
            
            # 判断是否超出当前分组区间
            if time_in_seconds >= start_time + interval:
                # 保存当前分组数据并开始新的分组
                grouped_data.append(current_group)
                current_group = []
                start_time = time_in_seconds  # 更新分组起始时间
            
            # 添加当前最大值到分组
            current_group.append(max_value)
        
        # 保存最后一组数据
        if current_group:
            grouped_data.append(current_group)
    
    # 保存每组数据到独立的文件
    # for i, group in enumerate(grouped_data):
    #     output_file = os.path.join(output_dir, f"group_{i + 1}.txt")
    #     with open(output_file, 'w') as out_file:
    #         out_file.write("\n".join(map(str, group)))
    #     print(f"已保存: {output_file}")
    return grouped_data


# 1 将音频数据分为 5秒一段 每小段如果包含鼾声，则认为发生打鼾
# annotations = pd.read_csv("E:\\data\\鼾声多模态\\止鼾枕标签\\止鼾枕清洗后的数据_out\\0102\\wang\\0102020000.csv")
# # 读取 气压文本
# press_data = process_log("E:\\data\\鼾声多模态\\止鼾枕清洗后的数据\\0102\\wang\\0102020000.log", output_files_path)

# wav_labels = generate_labels(30*60, annotations)
# split_audio("E:\\data\\鼾声多模态\\止鼾枕清洗后的数据\\0102\\wang\\0102020000_test.wav", output_files_path, wav_labels, press_data)


# 2 遍历所有的文件夹 进行操作

date_list = os.listdir(wav_files_path)  # 日期文件夹
print(date_list)
for date_index in date_list:
    name_path = os.path.join(wav_files_path, date_index)
    name_list = os.listdir(name_path)    # 姓名文件夹
    print(name_list)       
    for name_index in name_list:
        file_path = os.path.join(wav_files_path, date_index, name_index)
        file_name_list = os.listdir(file_path)
        for press_log in file_name_list:
            if press_log.split(".")[1] == 'log':
                press_data = process_log(os.path.join(file_path, press_log), output_files_path)
                annotations = pd.read_csv(os.path.join(wav_label_path, date_index, name_index, press_log.split(".")[0]+".csv"))
                # 获取音频时长
                temp_audio = AudioSegment.from_file(os.path.join(file_path, press_log.split(".")[0]+'_test.wav'))
                # 获取音频时长（毫秒）
                total_length = len(temp_audio)//1000 + 1  # 补1秒  
                wav_labels = generate_labels(total_length, annotations)
                file_nick_name = date_index +"_"+ name_index +"_"+ press_log.split(".")[0]
                split_audio(os.path.join(file_path, press_log.split(".")[0]+'_test.wav'), output_files_path, wav_labels, press_data, group_name=file_nick_name)
            # print(press_log)
        
